%% Sorting: firm Size and Local Density
% Calculate the Price Index at each point s

clc;
clear;
addpath Functions;
datadir = '..\..\Data\Final\TableX\';

P_s = zeros(36,1);

Qs = [3,4,5,7,9,10,12,13,15,17,18,20,22,23,25];
P_s_q = zeros(36,length(Qs));
y = length(Qs);
for m = 1:y
    q = Qs(m);
    load(strcat(datadir,'N36_sigma045_delta9_q',num2str(q),'.mat'));

    for loc = 1:length(plant_locs)
        loc_s = plant_locs(loc,:);
        loc_sum = 0;
        for plant_o = 1:length(plant_locs) 
            if choice(plant_o) == 1
                plant = plant_locs(plant_o,:);
                loc_sum = loc_sum + (BoW(plant_o) * Zprod(q,N_init*delta^2,sigma) / (((norm(loc_s-plant))/delta)+1))^theta;
            end
        end
        P_s_q(loc,m) = (epsilon/(epsilon - 1))*(loc_sum^(-1/theta));
    end
    
 
    
end

P_s_q = P_s_q.^(1-epsilon);

for k=1:length(P_s)
    P_s(k) = sum(P_s_q(k,:),2)^(1/(1-epsilon));
end
% Run this after running wrapper_square.m
Qsnew = [3,4,5,7,9,10,12,13,15,17,18,20,22,23,25];
len = length(Qsnew);
regressionmatrix = zeros(len,4);

ps_matrix = zeros(grid_sz);
idx = 1;
for i=6:-1:1
    for j = 1:6
        ps_matrix(i,j) = P_s(idx,1);
        idx=idx+1;
    end
end
ps_matrix = ps_matrix.^(-epsilon);
Ws = ones(grid_sz) ;


for m=1:len
    q = Qsnew(m);
    load(strcat(datadir,'N36_sigma045_delta9_q',num2str(q),'.mat'));
    
    Ls = (ps_matrix .* Ds(2:7,2:7)) ./ Ws;


    regressionmatrix(m,1) = log(q);
    regionalDemand = zeros(2,2);
    regionalDemand(1,1) = mean(Ls(1:3,1:3),'all');
    regionalDemand(1,2) = mean(Ls(1:3,4:6),'all');
    regionalDemand(2,1) = mean(Ls(4:6,1:3),'all');
    regionalDemand(2,2) = mean(Ls(4:6,4:6),'all');

    choicematrix = zeros(grid_sz);
    idx = 1;
    for i=6:-1:1
        for j = 1:6
            choicematrix(i,j) = choice(idx,1);
            idx=idx+1;
        end
    end

    regionalPlants = zeros(2,2);
    regionalPlants(1,1) = sum(choicematrix(1:3,1:3),'all');
    regionalPlants(1,2) = sum(choicematrix(1:3,4:6),'all');
    regionalPlants(2,1) = sum(choicematrix(4:6,1:3),'all');
    regionalPlants(2,2) = sum(choicematrix(4:6,4:6),'all');
    regionalPlants = regionalPlants ./ sum(regionalPlants,'all');

    weightedDemand = regionalPlants .* regionalDemand;
    weightedDemand(isinf(weightedDemand)) = 0;
    regressionmatrix(m,2) = log(sum(weightedDemand,'all'));
    regressionmatrix(m,3) = log(profits);
    regressionmatrix(m,4) = sum(choicematrix,'all');
end

cd '..\..\Data\Final\TableX\';
regression_tbl2 = array2table(regressionmatrix,...
    'VariableNames',{'log_q','log_weight_pop','log_tot_rev','plants'});
writetable(regression_tbl2,"Ls9_sorting_regression_tbl2.csv");

cd '..\..\..\Code\Setup\'
regressionmatrix2 = zeros(len,4);

for m=1:len
    q = Qsnew(m);
    load(strcat(datadir,'N36_sigma045_delta9_q',num2str(q),'.mat'));
    
    Ls = (ps_matrix .* Ds(2:7,2:7)) ./ Ws;
    
    regressionmatrix2(m,1) = log(q);
    regionalDemand = zeros(3,3);
    regionalDemand(1,1) = mean(Ls(1:2,1:2),'all');
    regionalDemand(1,2) = mean(Ls(1:2,3:4),'all');
    regionalDemand(1,3) = mean(Ls(1:2,5:6),'all');
    regionalDemand(2,1) = mean(Ls(3:4,1:2),'all');
    regionalDemand(2,2) = mean(Ls(3:4,3:4),'all');
    regionalDemand(2,3) = mean(Ls(3:4,5:6),'all');
    regionalDemand(3,1) = mean(Ls(5:6,1:2),'all');
    regionalDemand(3,2) = mean(Ls(5:6,3:4),'all');
    regionalDemand(3,3) = mean(Ls(5:6,5:6),'all');

    choicematrix = zeros(grid_sz);
    idx = 1;
    for i=6:-1:1
        for j = 1:6
            choicematrix(i,j) = choice(idx,1);
            idx=idx+1;
        end
    end

    regionalPlants = zeros(3,3);
    regionalPlants(1,1) = sum(choicematrix(1:2,1:2),'all');
    regionalPlants(1,2) = sum(choicematrix(1:2,3:4),'all');
    regionalPlants(1,3) = sum(choicematrix(1:2,5:6),'all');
    regionalPlants(2,1) = sum(choicematrix(3:4,1:2),'all');
    regionalPlants(2,2) = sum(choicematrix(3:4,3:4),'all');
    regionalPlants(2,3) = sum(choicematrix(3:4,5:6),'all');
    regionalPlants(3,1) = sum(choicematrix(5:6,1:2),'all');
    regionalPlants(3,2) = sum(choicematrix(5:6,3:4),'all');
    regionalPlants(3,3) = sum(choicematrix(5:6,5:6),'all');
    regionalPlants = regionalPlants ./ sum(regionalPlants,'all');

    weightedDemand = regionalPlants .* regionalDemand;
    weightedDemand(isinf(weightedDemand)) = 0;
    regressionmatrix2(m,2) = log(sum(weightedDemand,'all'));
    regressionmatrix2(m,3) = log(profits);
    regressionmatrix2(m,4) = sum(choicematrix,'all');
end

cd '..\..\Data\Final\TableX\';
regression_tbl1 = array2table(regressionmatrix2,...
    'VariableNames',{'log_q','log_weight_pop','log_tot_rev','plants'});
writetable(regression_tbl1,"Ls9_sorting_regression_tbl1.csv");